from PySide2.QtWidgets import QApplication, QWidget
from PySide2.QtCore import Signal, Qt, QPoint, QRect
from PySide2.QtGui import QPalette, QBrush, QPainter, QPen, QColor
import pytesseract
from PIL import ImageGrab
import sys


class CaptureWidget(QWidget):

    # 自定义信号
    send_str_signal = Signal(str)  # 将结果传给主界面
    close_capture_signal = Signal()

    def __init__(self, lang):
        super().__init__()

        self._lang = lang

        self.setWindowTitle('Capturing')
        self.setWindowFlags(Qt.FramelessWindowHint  # 无边框
                            | Qt.WindowStaysOnTopHint  # 页面置于最顶层
                            | Qt.Dialog)  # dialog样式，无按钮
        self.setWindowState(Qt.WindowFullScreen)  # 全屏

        # 获取当前屏幕的图像
        self.screen = QApplication.primaryScreen().grabWindow(0)

        # 配置调色板，将当前屏幕的图像画在widget上
        palette = QPalette()
        palette.setBrush(self.backgroundRole(), QBrush(self.screen))
        self.setPalette(palette)

        # QApplication.setOverrideCursor(Qt.CrossCursor)  # 设置全局的鼠标样式
        self.setCursor(Qt.CrossCursor)  # 设置鼠标样式为十字
        self.start, self.end = QPoint(), QPoint()  # 初始化起始点和终点

    def keyPressEvent(self, event):
        if event.key() == Qt.Key_Escape:  # esc键退出
            self.close_capture_signal.emit()  # 关闭窗口

        return super().keyPressEvent(event)

    def paintEvent(self, event):
        if self.start == self.end:
            return super().paintEvent(event)

        # 设置绘图方式
        painter = QPainter(self)
        painter.setPen(QPen(QColor(255, 0, 0), 3))  # 画笔颜色，粗细
        painter.drawRect(QRect(self.start, self.end))
        return super().paintEvent(event)

    def mousePressEvent(self, event):
        self.start = self.end = event.pos()
        self.update()
        return super().mousePressEvent(event)

    def mouseMoveEvent(self, event):
        self.end = event.pos()
        self.update()
        return super().mousePressEvent(event)

    def mouseReleaseEvent(self, event):
        if self.start == self.end:
            return super().mouseReleaseEvent(event)

        x1, x2 = sorted((self.start.x(), self.end.x()))
        y1, y2 = sorted((self.start.y(), self.end.y()))
        screenShot = ImageGrab.grab(bbox=(x1, y1, x2, y2))
        ocr_str = self.img2str(screenShot)

        # 将识别的结果发送出去
        self.send_str_signal.emit(ocr_str)

        # 处理完关闭当前widget
        self.close_capture_signal.emit()

    # 解析图片
    def img2str(self, img):
        return pytesseract.image_to_string(img, lang=self._lang,
                                           timeout='5')  # lang语言选择为中文


def main():
    app = QApplication(sys.argv)
    mainw = CaptureWidget('chi_sim')
    mainw.show()
    sys.exit(app.exec_())


if __name__ == '__main__':
    main()
